from flask import Flask, render_template, request, jsonify, send_from_directory, Response
import sys
import os
import psycopg2
import ast
from jinja2 import Environment

from agent import functions

# 让 Flask 能找到 agent 包
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from agent.engine import get_env, render_template, render_template_stream

app = Flask(__name__)

PG_CONFIG = {
    'host': '192.168.8.187',
    'port': 5432,
    'database': 'agent',
    'user': 'tianye',
    'password': 'a1991524'
}

def is_safe_function_code(code):
    try:
        tree = ast.parse(code)
        for node in ast.walk(tree):
            if isinstance(node, (ast.Import, ast.ImportFrom, ast.Global, ast.Nonlocal, ast.With, ast.Try, ast.Raise)):
                return False
            if isinstance(node, ast.Name):
                if node.id.startswith('__') or node.id in {'exec', 'eval', 'open', 'os', 'sys', 'subprocess'}:
                    return False
        return True
    except Exception:
        return False

def get_conn():
    return psycopg2.connect(**PG_CONFIG)

@app.route('/users', methods=['GET', 'POST'])
def users():
    if request.method == 'GET':
        with get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT id, username FROM users ORDER BY id")
                users = [{'id': row[0], 'username': row[1]} for row in cur.fetchall()]
        return jsonify(users)
    else:
        username = request.json.get('username')
        if not username:
            return jsonify({'error': '用户名不能为空'}), 400
        with get_conn() as conn:
            with conn.cursor() as cur:
                try:
                    cur.execute("INSERT INTO users (username) VALUES (%s) RETURNING id", (username,))
                    user_id = cur.fetchone()[0]
                    conn.commit()
                    return jsonify({'id': user_id, 'username': username})
                except Exception as e:
                    return jsonify({'error': str(e)}), 400

@app.route('/user_functions', methods=['GET', 'POST', 'DELETE'])
def user_functions():
    user_id = request.args.get('user_id') or request.json.get('user_id')
    if not user_id:
        return jsonify({'error': 'user_id 必填'}), 400
    if request.method == 'GET':
        with get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute("SELECT id, func_name, func_code FROM user_functions WHERE user_id = %s ORDER BY id", (user_id,))
                funcs = [{'id': row[0], 'func_name': row[1], 'func_code': row[2]} for row in cur.fetchall()]
        return jsonify(funcs)
    elif request.method == 'POST':
        func_id = request.json.get('id')
        update = request.json.get('update')
        func_name = request.json.get('func_name')
        func_code = request.json.get('func_code')
        if update and func_id:
            # 更新
            if not func_code:
                return jsonify({'error': '函数代码不能为空'}), 400
            if not is_safe_function_code(func_code):
                return jsonify({'error': '函数代码不安全'}), 400
            with get_conn() as conn:
                with conn.cursor() as cur:
                    cur.execute("UPDATE user_functions SET func_code=%s WHERE id=%s", (func_code, func_id))
                    conn.commit()
            return jsonify({'id': func_id, 'func_name': func_name, 'func_code': func_code})
        # 新增
        if not func_name or not func_code:
            return jsonify({'error': '函数名和代码不能为空'}), 400
        if not is_safe_function_code(func_code):
            return jsonify({'error': '函数代码不安全'}), 400
        with get_conn() as conn:
            with conn.cursor() as cur:
                try:
                    cur.execute("INSERT INTO user_functions (user_id, func_name, func_code) VALUES (%s, %s, %s) RETURNING id", (user_id, func_name, func_code))
                    func_id = cur.fetchone()[0]
                    conn.commit()
                    return jsonify({'id': func_id, 'func_name': func_name, 'func_code': func_code})
                except Exception as e:
                    return jsonify({'error': str(e)}), 400
    elif request.method == 'DELETE':
        func_id = request.json.get('id')
        if not func_id:
            return jsonify({'error': 'id 必填'}), 400
        with get_conn() as conn:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM user_functions WHERE id = %s", (func_id,))
                conn.commit()
        return jsonify({'success': True})

@app.route('/run', methods=['POST'])
def run_template():
    data = request.json
    template_code = data.get('template', '')
    user_id = data.get('user_id')
    question = data.get('question')
    history = data.get('history', [])
    session_id = data.get('session_id')
    if not user_id:
        return jsonify({'result': 'user_id 必填'}), 400
    # 查询并注册用户自定义函数
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT func_name, func_code FROM user_functions WHERE user_id = %s", (user_id,))
            user_funcs = cur.fetchall()
    user_func_dict = {}
    for func_name, func_code in user_funcs:
        local_ns = {}
        try:
            exec(func_code, {}, local_ns)
            if func_name in local_ns:
                user_func_dict[func_name] = local_ns[func_name]
        except Exception as e:
            print(f"注册函数 {func_name} 失败: {e}")
    context = {'question': question, 'history': history, 'session_id': session_id}
    result = render_template(template_code, context)
    print("RESULT:", result)
    return jsonify({'result': result})

@app.route('/run_stream', methods=['POST'])
def run_template_stream():
    data = request.json
    template_code = data.get('template', '')
    user_id = data.get('user_id')
    question = data.get('question')
    history = data.get('history', [])
    session_id = data.get('session_id')
    if not user_id:
        return Response('user_id 必填', status=400)
    # 查询并注册用户自定义函数
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT func_name, func_code FROM user_functions WHERE user_id = %s", (user_id,))
            user_funcs = cur.fetchall()
    user_func_dict = {}
    for func_name, func_code in user_funcs:
        local_ns = {}
        try:
            exec(func_code, {}, local_ns)
            if func_name in local_ns:
                user_func_dict[func_name] = local_ns[func_name]
        except Exception as e:
            print(f"注册函数 {func_name} 失败: {e}")
    context = {'question': question, 'history': history, 'session_id': session_id}
    def generate():
        for chunk in render_template_stream(template_code, context):
            yield f"data: {chunk}\n\n"
    return Response(generate(), mimetype='text/event-stream')

@app.route('/test_jinja', methods=['POST'])
def test_jinja():
    data = request.json
    template_code = data.get('template', '')
    env = Environment()
    template = env.from_string(template_code)
    result = template.render()
    print("TEST_JINJA RESULT:", result)
    return jsonify({'result': result})

@app.route('/save_template', methods=['POST'])
def save_template():
    data = request.json
    user_id = data.get('user_id')
    name = data.get('name')
    content = data.get('content')
    if not user_id or not name or not content:
        return jsonify({'success': False, 'error': '参数缺失'})
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(
                "INSERT INTO user_templates (user_id, name, content) VALUES (%s, %s, %s)",
                (user_id, name, content)
            )
            conn.commit()
    return jsonify({'success': True})

@app.route('/user_templates')
def user_templates():
    user_id = request.args.get('user_id')
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT id, name FROM user_templates WHERE user_id=%s ORDER BY created_at DESC", (user_id,))
            templates = [{'id': row[0], 'name': row[1]} for row in cur.fetchall()]
    return jsonify(templates)

@app.route('/load_template')
def load_template():
    tid = request.args.get('id')
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT content FROM user_templates WHERE id=%s", (tid,))
            row = cur.fetchone()
    return jsonify({'content': row[0] if row else ''})

@app.route('/ai_generate_template', methods=['POST'])
def ai_generate_template():
    data = request.json
    user_id = data.get('user_id')
    user_prompt = data.get('prompt')
    # 查询当前用户所有函数
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute("SELECT func_name, func_code FROM user_functions WHERE user_id=%s", (user_id,))
            user_functions = cur.fetchall()
    # 优化后的系统提示词
    sys_prompt = (
        "你是一个Jinja2智能模板生成助手。\n"
        "当前系统已有的函数如下：\n"
    )
    for fname, fcode in user_functions:
        sys_prompt += f"函数名: {fname}\n代码:\n{fcode}\n\n"
    sys_prompt += (
        "请根据用户的需求，生成一个Jinja2模板，模板要能直接用于AI对话流程。\n"
        "如果需要调用大模型，请直接用 {% set result = gen(question) %}，其中 question 是用户输入的问题。\n"
        "如果需要构建复杂的提示词，必须用 Jinja2 的 set 块，并且一定要把用户输入（question 变量）拼接到 prompt 里，例如：\n"
        "{% set prompt %}\n请识别用户的意图，并将其分类为以下三种之一：打车、坐船、坐飞机。\n用户输入：{{ question }}\n请在识别后打印固定的输出。\n{% endset %}\n"
        "然后用 {% set result = gen(prompt) %} 获取大模型回复后在继续下面的逻辑。\n"
        "模板中只输出 Jinja2 代码，不要输出任何解释、HTML、markdown代码块标记（如 ```jinja2 或 ```），只要纯模板内容。\n"
        "用户需求：" + (user_prompt or '')
    )
    # 调用大模型
    template = functions.gen(sys_prompt, model='Qwen/Qwen2.5-Coder-7B-Instruct', max_tokens=2048)
    return jsonify({'template': template})

@app.route('/')
def index():
    return send_from_directory('static', 'index.html')

if __name__ == '__main__':
    app.run(debug=True)
